#!/usr/bin/env python3
"""
简单使用示例：用 datasets.load_dataset 读取 parquet 文件
"""

from datasets import load_dataset
import json

def main():
    # 加载训练数据
    print("加载训练数据...")
    dataset = load_dataset("parquet", data_files="./g3_hf_dataset/train.parquet")

    # 获取训练集
    train_data = dataset["train"]
    print(f"训练数据量: {len(train_data)}")

    # 查看第一条数据
    sample = train_data[0]

    print(f"\n数据样例:")
    print(f"查询ID: {sample['query_id']}")
    print(f"查询: {sample['query']}")
    print(f"查询Agent_name: {sample['agent_name']}")
    print(f"task : {sample['task']}")
    print(f"可用API: {sample['api_names']}")
    print(f"相关API: {sample['relevant_apis']}")

    # 获取原始API数据
    original_apis = json.loads(sample['api_list_ori'])
    print(f"\n原始API数量: {len(original_apis)}")
    print(f"第一个API: {original_apis[0]['api_name']} - {original_apis[0]['tool_name']}")
    print(sample.keys())

if __name__ == "__main__":
    main()